import pandas as pd

from app_config import get_engine_ts, get_pro
import numpy as np
import time

# 一年前
one_year_ago = '20231101'
current_date_str = '20241030'
one_season_ago = "20240801"


def get_profit(_ts_code):
    for _ in range(10):
        try:
            df_ = get_pro().query('fina_indicator', ts_code=_ts_code,
                                  fields="ts_code,end_date,q_opincome,q_investincome,q_dtprofit,q_eps,q_netprofit_margin,"
                                         "q_gsprofit_margin,q_exp_to_sales,q_profit_to_gr,q_saleexp_to_gr,q_adminexp_to_gr,"
                                         "q_finaexp_to_gr,q_impair_to_gr_ttm,q_gc_to_gr,q_op_to_gr,q_roe,q_dt_roe,q_npta,"
                                         "q_opincome_to_ebt,q_investincome_to_ebt,q_dtprofit_to_profit,q_salescash_to_or,"
                                         "q_ocf_to_sales,q_ocf_to_or", start_date=one_year_ago,
                                  end_date=current_date_str)

            df_['q_profit'] = df_['q_dtprofit'] / df_['q_dtprofit_to_profit'] * 100
            print("财务数据 " + _ts_code + "  :  " + str(len(df_)))
            if len(df_) > 4:
                print("err =================" + str(len(df_)))
            df_dedup = df_.drop_duplicates(subset='end_date', keep='first')
            total_q_dtprofit = df_dedup['q_dtprofit'].sum()
            total_q_profit = df_dedup['q_profit'].sum()
            return total_q_dtprofit / 100000000, total_q_profit / 100000000
        except Exception as e:
            print(f"An unexpected error occurred: {e}")
            time.sleep(10)
    print("=========================================")
    print("=========================================")
    print("error error error error error" + _ts_code)
    print("=========================================")
    print("=========================================")
    return 0.0, 0.0


def calculate_volatility(df):
    df = df.sort_values(by='trade_date')

    df_one_year = df.copy()  # 使用.copy()显式创建副本



    # 计算每日对数收益率
    df_one_year['log_return'] = np.log(df_one_year['close'] / df_one_year['close'].shift(1))

    # 删除缺失值
    df_one_year = df_one_year.dropna(subset=['log_return'])

    # 计算波动率（年化）
    volatility = df_one_year['log_return'].std() * np.sqrt(len(df_one_year))  # 252是交易日的近似值

    return volatility


if __name__ == '__main__':
    from _daily_basic import fetch_data as fetch_data

    fetch_data(one_year_ago, current_date_str)
    from _stock_basic import get_data as get_stock_basic

    engine = get_engine_ts()

    # # 获取今天的日期
    # today = datetime.now()
    """     1 样本空间    """
    ''' 1.1 上市时间超过 12 个月； '''
    stock_basic = get_stock_basic()

    print("全部股票: " + str(len(stock_basic)))

    ''' 1.2 被实施退市风险警示除外； '''
    # 删除 name 列包含 'ST' 字符串的行
    st_filtered = stock_basic[~stock_basic['name'].str.contains('ST')]

    print("排除ST股票: " + str(len(st_filtered)))

    # 根据条件过滤数据
    filtered_df = st_filtered[(
                                      (st_filtered['market'].isin(['科创板', '创业板']))
                                      &
                                      (st_filtered['list_date'] <= one_year_ago)
                              ) | (
                                      (st_filtered['market'] == '主板')
                                      &
                                      (st_filtered['list_date'] <= one_season_ago)
                              )]
    print("符合上市时间要求股票: " + str(len(filtered_df)))

    """     2、选样方法    """
    '''2.1对样本空间内的证券按照过去一年的日均成交金额由高到低排名，剔除排名后 10%的证券作为待选样本；'''
    # 转换 ts_code 列为元组，并生成 SQL 中 IN 子句所需的格式
    ts_codes = filtered_df['ts_code'].tolist()

    # 构建查询语句
    query = f"""
    SELECT * FROM `daily_basic`
    WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
    AND trade_date >= '{one_year_ago}'
    AND trade_date <= '{current_date_str}'
    """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)

    '''circ_mv : 流通市值 * turnover_rate: 换手率  = 成交额'''
    result_df['amount'] = result_df['circ_mv'] * result_df['turnover_rate']

    '''对待选样本按照过去一年的日均总市值由高到低排名，选取排名前 50的证券作为指数样本。'''

    # 以 ts_code 列分组，并计算 total_mv 列的平均值
    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', 'avg_amount']
    merge = pd.merge(filtered_df, grouped_df, on='ts_code', how='left')

    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', 'avg_total_mv']
    pd_merge = pd.merge(merge, grouped_df_total_mv, on='ts_code', how='left')

    # 按 avg_amount 列降序排序
    df_sorted = pd_merge.sort_values(by='avg_amount', ascending=False)
    # 获取前 50% 的数据
    num_rows = len(df_sorted)
    top_50_percent = df_sorted.head(int(num_rows * 0.5))

    top_50_percent = top_50_percent.sort_values('avg_total_mv', ascending=False)
    top_50_percent = top_50_percent.head(300)

    d = top_50_percent

    '''计算TTM'''
    # 使用 apply 方法添加两列 total_q_dtprofit 和 total_q_profit
    # todo
    d[['扣非净利润(总)', '净利润(总)']] = d['ts_code'].apply(lambda _ts_code: pd.Series(get_profit(_ts_code)))

    data_frame_d = d.reset_index()

    data_frame_d.to_excel('top_300_备选.xlsx', sheet_name="bei")
